# web/server.py

import sys
import os
import re
from flask import Flask, request, jsonify, render_template
from flask_cors import CORS

# 将父目录（项目根目录）添加到Python路径中
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from pysql_compiler.compiler import Compiler
from pysql_compiler.catalog import SemanticError
from pysql_compiler.lexer import Lexer  # <<< 新增导入
from engine.executor import Executor
from engine.catalog_manager import CatalogManager
from engine.storage_engine import StorageEngine
from storage.buffer import BufferManager, LRUReplacementPolicy
from storage.file_manager import FileManager

app = Flask(__name__, template_folder='.', static_folder='.')
CORS(app)

# --- 核心修正：确保数据库文件路径的绝对和统一 ---
# 构造到项目根目录的绝对路径
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
db_file = os.path.join(project_root, "mydb.db")

print(f"--- [Server] 正在连接数据库文件: {db_file} ---")

file_manager = FileManager(db_file)
buffer_manager = BufferManager(pool_size=10, file_manager=file_manager, replacement_policy_class=LRUReplacementPolicy)
catalog_manager = CatalogManager(db_file)
storage_engine = StorageEngine(buffer_manager, catalog_manager)
compiler = Compiler(catalog=catalog_manager._catalog)
executor = Executor(catalog_manager=catalog_manager, storage_engine=storage_engine, compiler=compiler)

print("--- 数据库引擎已在服务器后端初始化 ---")


# --- 辅助函数 (无需修改) ---

def _split_sql_statements(sql_script: str) -> list[str]:
    """智能分割SQL脚本，并移除 '--' 风格的注释。"""
    script_no_comments = []
    for line in sql_script.splitlines():
        script_no_comments.append(line.split('--')[0])
    cleaned_script = "\n".join(script_no_comments)
    statements = re.split(r";(?=(?:[^']*'[^']*')*[^']*$)", cleaned_script)
    return [stmt.strip() for stmt in statements if stmt.strip()]


def _format_results_for_log(header: list, rows: list) -> str:
    """将查询结果格式化为对齐的文本表格，用于日志显示。"""
    if not rows:
        return "\n(0 rows returned)"

    if not header:  # 如果没有表头，直接返回行数
        return f"\n({len(rows)} rows returned, no header info)"

    widths = [len(str(h)) for h in header]
    for row in rows:
        for i, cell in enumerate(row):
            if i < len(widths):
                widths[i] = max(widths[i], len(str(cell)))

    separator = '-+-'.join('-' * w for w in widths)
    header_line = ' | '.join(str(h).ljust(w) for h, w in zip(header, widths))
    row_lines = [' | '.join(str(cell).ljust(w) for cell, w in zip(row, widths)) for row in rows]

    return '\n' + '\n'.join([header_line, separator] + row_lines)


def find_project_node(p):
    if p and isinstance(p, dict) and p.get('operation') == 'PROJECT': return p
    if p and isinstance(p, dict) and 'source' in p: return find_project_node(p.get('source'))
    return None


def extract_header_from_plan(plan, catalog_manager):
    project_plan = find_project_node(plan)
    if not project_plan: return []
    header = []
    is_star = project_plan.get('columns') and project_plan['columns'][0]['type'] == 'star'
    if is_star:
        source_plan = project_plan.get('source', {})
        table_name_node = source_plan
        while isinstance(table_name_node, dict) and 'source' in table_name_node:
            table_name_node = table_name_node['source']
        table_name = table_name_node.get('table_name')
        if table_name:
            schema = catalog_manager.get_table_schema(table_name)
            if schema: return schema.column_order
    else:
        for col in project_plan.get('columns', []):
            if col.get('type') == 'alias':
                header.append(col['alias']['name'])
            elif col.get('type') == 'identifier':
                header.append(col['name'])
    return header


# --- API 路由 (部分需要修改) ---

@app.route('/')
def index():
    return render_template('index.html')


@app.route('/api/schema', methods=['GET'])
def get_schema_info():
    try:
        schema_data = {"tables": []}
        table_names = catalog_manager.get_all_table_names()
        for table_name in table_names:
            schema = catalog_manager.get_table_schema(table_name)
            columns = []
            if schema and hasattr(schema, 'columns'):
                for col_name in schema.column_order:
                    data_type = schema.columns.get(col_name, 'UNKNOWN')
                    columns.append({"name": col_name, "type": str(data_type), "is_primary_key": False})
            schema_data["tables"].append({"name": table_name, "columns": columns})
        return jsonify({"success": True, "schema": schema_data})
    except Exception as e:
        return jsonify({"success": False, "error": f"[服务器错误] 无法获取Schema: {e}"}), 500


@app.route('/api/table_data/<table_name>', methods=['GET'])
def get_table_data(table_name):
    try:
        page = request.args.get('page', 1, type=int)
        limit = request.args.get('limit', 50, type=int)
        offset = (page - 1) * limit
        data_sql = f"SELECT * FROM {table_name};"
        plan_tuple = compiler.compile(data_sql)
        actual_plan = plan_tuple[0] if isinstance(plan_tuple, tuple) else plan_tuple
        all_rows = executor.execute(actual_plan) or []
        total_rows = len(all_rows)
        rows = all_rows[offset:offset + limit]
        header = []
        schema = catalog_manager.get_table_schema(table_name)
        if schema: header = schema.column_order
        return jsonify({
            "success": True, "table_name": table_name, "header": header, "rows": rows,
            "total_rows": total_rows, "page": page, "limit": limit
        })
    except Exception as e:
        return jsonify({"success": False, "error": f"获取表 '{table_name}' 数据失败: {e}"}), 500


@app.route('/api/execute_sql', methods=['POST'])
def execute_sql_endpoint():
    sql_query = request.json.get('sql', '').replace('\r\n', '\n')
    if not sql_query:
        return jsonify({"error": "SQL query is empty"}), 400

    statements = _split_sql_statements(sql_query)
    if not statements:
        return jsonify({"success": True, "message": "No executable statements found.", "header": [], "rows": [],
                        "schema_modified": False})

    results = []
    final_header = []
    schema_modified = False
    last_plan_with_results = None
    response_data = {}
    statement_to_report_error = ""

    try:
        for statement in statements:
            statement_to_report_error = statement  # 记录当前语句以备报错时使用
            plan_tuple = compiler.compile(statement + ';')
            actual_plan = plan_tuple[0] if isinstance(plan_tuple, tuple) else plan_tuple
            plans_to_run = actual_plan if isinstance(actual_plan, list) else [actual_plan]

            for p in plans_to_run:
                op_type = p.get('operation', '')
                if op_type in ['CREATE_TABLE', 'DROP_TABLE', 'ALTER_TABLE', 'INSERT', 'DELETE', 'UPDATE']:
                    schema_modified = True

                res = executor.execute(p)

                if res is not None:
                    last_plan_with_results = p
                    if isinstance(res, list):
                        results = res
                    else:
                        results = [res]

        if last_plan_with_results:
            final_header = extract_header_from_plan(last_plan_with_results, catalog_manager)
            message = f"查询成功，返回 {len(results)} 行数据。"
        else:
            message = f"成功执行 {len(statements)} 条语句。"

        response_data = {
            "success": True,
            "message": message,
            "header": final_header,
            "rows": results,
            "schema_modified": schema_modified
        }
        return jsonify(response_data)

    except Exception as e:
        return jsonify(
            {"success": False, "error": f"[数据库错误] 在执行 '{statement_to_report_error}' 时发生错误: {e}"}), 500

    finally:
        if schema_modified:
            print("--- [Server] 检测到数据/结构变更，正在将缓存刷入磁盘... ---")
            buffer_manager.flush_all_pages()
            catalog_manager.save_catalog()
            print("--- [Server] 刷新完成。 ---")


# <<<--- 新增API端点，用于语法补全 ---<<<
@app.route('/api/suggestions', methods=['POST'])
def get_suggestions():
    data = request.json
    query_so_far = data.get('query', '').upper()

    suggestions = []

    # 1. 从 Lexer 中获取所有定义的关键字
    lexer_keywords = list(Lexer("").keywords)

    # 2. 从 CatalogManager 获取所有的表名
    table_names = catalog_manager.get_all_table_names()

    # 3. 从 CatalogManager 获取所有表的列名
    all_columns = set()
    for table_name in table_names:
        schema = catalog_manager.get_table_schema(table_name)
        if schema:
            all_columns.update(schema.column_order)

    # 4. 基于简化的上下文分析，提供建议
    words = query_so_far.split()
    last_word = words[-1] if words and query_so_far.endswith(' ') else ''

    if not words or (len(words) == 1 and not query_so_far.endswith(' ')):
        # 查询开头：建议所有关键字和表名
        suggestions.extend(lexer_keywords)
        suggestions.extend(table_names)
    elif last_word in ['FROM', 'JOIN', 'UPDATE', 'INTO', 'TABLE']:
        # 在需要表名的关键字后：建议所有表名
        suggestions.extend(table_names)
    elif last_word in ['WHERE', 'ON', 'BY', 'SET', '(', ',']:
        # 在需要列名的关键字后：建议所有列名
        suggestions.extend(list(all_columns))
    else:
        # 默认情况：建议所有内容，让前端去过滤
        suggestions.extend(lexer_keywords)
        suggestions.extend(table_names)
        suggestions.extend(list(all_columns))

    # 5. 去重并排序
    final_suggestions = sorted(list(set(s.lower() for s in suggestions)))

    return jsonify({"success": True, "suggestions": final_suggestions})


# >>>--- 新增结束 ---<<<


@app.route('/api/execute_script', methods=['POST'])
def execute_script_endpoint():
    sql_script = request.json.get('sql', '').replace('\r\n', '\n')
    if not sql_script: return jsonify({"success": False, "error": "SQL script is empty"}), 400

    statements = _split_sql_statements(sql_script)

    log_messages = []
    success_count = 0
    error_count = 0
    schema_modified_in_script = False
    total_statements = len(statements)

    for statement in statements:
        try:
            plan_tuple = compiler.compile(statement + ';')
            actual_plan = plan_tuple[0]

            plans_to_run = actual_plan if isinstance(actual_plan, list) else [actual_plan]

            final_result = None
            op_type = ''

            for plan in plans_to_run:
                if not isinstance(plan, dict):
                    raise TypeError("执行计划必须是一个字典。")

                op_type = plan.get('operation', '')
                if op_type in ['CREATE_TABLE', 'DROP_TABLE', 'ALTER_TABLE', 'INSERT', 'DELETE', 'UPDATE']:
                    schema_modified_in_script = True

                result = executor.execute(plan)
                if result is not None:
                    final_result = result

            success_count += 1
            log_entry = f"[成功] >> {statement}"

            if final_result and op_type == 'PROJECT':
                header = extract_header_from_plan(plans_to_run[0], catalog_manager)
                log_entry += _format_results_for_log(header, final_result)

            log_messages.append(log_entry)

        except Exception as e:
            error_count += 1
            log_messages.append(f"[失败] >> {statement}\n     错误原因: {e}")

    if schema_modified_in_script:
        print("--- [Server] 脚本执行完毕，检测到数据/结构变更，正在将缓存刷入磁盘... ---")
        buffer_manager.flush_all_pages()
        catalog_manager.save_catalog()
        print("--- [Server] 刷新完成。 ---")

    summary = {
        "total_statements": total_statements, "success_count": success_count, "error_count": error_count,
        "message": f"脚本执行完毕。总共 {total_statements} 条语句，成功 {success_count} 条，失败 {error_count} 条。",
        "schema_modified": schema_modified_in_script
    }

    return jsonify({"success": True, "summary": summary, "log": log_messages})


if __name__ == '__main__':
    app.run(debug=True, port=5000)
